
"""
@author: xiangping
@contact: xiangpingbu@gmail.com
@time: 2020/1/5 4:15 下午
@file: FileModel
@Desc:
"""
import joblib
import json
from app.configs.model_path import MODEL_DICT
from app.services.models.base_model import BaseModel
import pickle


class FileModel(BaseModel):
    def load_model(self):
        if ".pkl" in MODEL_DICT[self.model_name]:
            return pickle.load(open(MODEL_DICT[self.model_name], "rb"))
        else:
            return joblib.load(MODEL_DICT[self.model_name])